import matplotlib.pyplot as plt
import numpy as np

from collections import namedtuple
import random

Transition = namedtuple('Transition',('state','action','next_state','reward'))
    
class ReplayMemory:
    def __init__(self,capacity):
        self.capacity = capacity
        self.memory = []
        self.index = 0
    def push(self,state,action,next_state,reward):
        if len(self.memory) < self.capacity:
            self.memory.append(None)
        self.memory[self.index] = Transition(state,action,next_state,reward)
        self.index = (self.index+1)%self.capacity
    def sample(self,batch_size):
        return random.sample(self.memory,batch_size)
    def __len__(self):
        return len(self.memory)

def huber_loss(delta,beta=1):
    if np.abs(delta) < beta:
        return 0.5*delta**2/beta
    return abs(delta) - 0.5*beta
def square_loss(delta):
    return 0.5*delta**2

# deltas = np.arange(-5,5,0.01)
# plt.plot(deltas,[huber_loss(delta) for delta in deltas])
# plt.plot(deltas,square_loss(deltas))
# plt.xticks(np.arange(-5,5,step=1))
# plt.legend(['huber loss','square loss'])

from torch import nn
import torch
import torch.nn.functional as F
from torch import optim

class DQN(nn.Module):
    def __init__(self,n_states,n_actions):
        super(DQN,self).__init__()
        self.fc1 = nn.Linear(n_states,128)
        self.fc2 = nn.Linear(128,128)
        self.fc3 = nn.Linear(128,n_actions)
    def forward(self,x):
        return self.fc3(F.relu(self.fc2(F.relu(self.fc1(x)))))
    
# t = torch.rand(3,2)
# print(t)
# print(t.max(dim=1)[0])

class Agent:
    def __init__(self,n_states,n_actions,eta=0.5,gamma=0.99,capacity=10000,batch_size=32):
        self.n_states=n_states
        self.n_actions=n_actions
        self.eta=eta
        self.gamma=gamma
        self.batch_size=batch_size

        self.memory = ReplayMemory(capacity)
        self.model = DQN(n_states,n_actions)
        self.optimizer = optim.Adam(self.model.parameters(),lr=0.0001)
    
    def _replay(self):
        if len(self.memory) < self.batch_size:
            return
        batch = self.memory.sample(self.batch_size)
        batch = Transition(*zip(*batch))
        state_batch = torch.cat(batch.state)
        action_batch = torch.cat(batch.action)
        reward_batch = torch.cat(batch.reward)

        non_final_next_state_batch = torch.cat([s for s in batch.next_state if s is not None])

        self.model.eval()

        state_action_values = self.model(state_batch).gather(dim=1,index=action_batch)
        non_final_mask = torch.ByteTensor(tuple(map(lambda s:s is not None,batch.next_state)))
        next_state_values = torch.zeros(self.batch_size)
        next_state_values[non_final_mask] = self.model(non_final_next_state_batch).max(dim=1)[0].detach()
        expected_state_action_values = reward_batch+self.gamma*next_state_values

        self.model.train()

        loss = F.smooth_l1_loss(state_action_values,expected_state_action_values.unsqueeze(1))
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
    
    def update_q_function(self):
        self._replay()
    
    def memorize(self,state,action,next_state,reward):
        self.memory.push(state,action,next_state,reward)
    
    def choose_action(self,state,episode):
        eps = 0.5*1/(1+episode)
        if random.random() <eps:
            action = torch.IntTensor([[random.randrange(self.n_actions)]])
        else:
            self.model.eval()
            with torch.no_grad():
                action = self.model(state).max(1)[1].view(1,1)
        return action

import gym
from matplotlib import animation

env = gym.make('CartPole-v0',render_mode = "rgb_array")
n_states = env.observation_space.shape[0]
n_actions = env.action_space.n

max_episodes = 500
max_step = 200

complete_episodes = 0
finished_flag = False

agent = Agent(n_states,n_actions)
frames = [] 
for episode in range(max_episodes):
    state = env.reset()[0] 
    state = torch.from_numpy(state).type(torch.FloatTensor).unsqueeze(0)
    for step in range(max_step):
        if finished_flag == True:
            frames.append(env.render())
        action = agent.choose_action(state,episode)

        next_state,_,done,info,_ = env.step(action.item())

        if done:
            next_state = None
            if step < 100:
                reward = torch.FloatTensor([-1,])
                complete_episodes = 0
            else:
                reward = torch.FloatTensor([1,])
                complete_episodes+=1
        else:
            reward = torch.FloatTensor([0])
            next_state = torch.from_numpy(next_state).type(torch.FloatTensor)
            next_state = next_state.unsqueeze(0)
        agent.memorize(state,action,next_state,reward)
        agent.update_q_function()
        state = next_state

        if done:
            print(f'episode:{episode},steps:{step}')
            break
    if finished_flag == True:
        break
    if complete_episodes >= 10:
        finished_flag = True
        print('连续成功10轮')
        
def display_frames_as_gif(frames,output):
    fig = plt.figure(figsize=(frames[0].shape[1]/72.0,frames[0].shape[0]/72.0),dpi=72)
    patch = plt.imshow(frames[0])
    plt.axis('off')

    def animate(i):
        img = patch.set_data(frames[i]) 
        return img
    anim = animation.FuncAnimation(plt.gcf(),animate,frames=len(frames),interval=50)
    anim.save(output)

if len(frames)>0:
    display_frames_as_gif(frames, 'E:/chartGPT_app/Reinforcement_learning/dqn_cartpole.gif')







